tensorrt_llm/_torch/auto_deploy/export/export.py (1)

280-280: Constant re-created on every loop iteration.

_tensor_list_types is invariant — hoist it above the for node in list(graph.nodes) loop (e.g., next to moe_ops).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/auto_deploy/export/export.py` at line 280, The tuple
_tensor_list_types is being recreated inside the loop; move its definition out
of the loop so it is created once (hoist it to the module scope near the
existing moe_ops declaration) and then reference _tensor_list_types inside the
for node in list(graph.nodes) loop to avoid repeated allocations.

tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py (2)

331-335: Nested _count_moe_experts duplicates the module-level _count_moe_experts_in_graph.

These two functions are identical. Now that _count_moe_experts_in_graph lives at module scope (lines 394-399), the nested copy can be eliminated.

♻️ Proposed fix

-    def _count_moe_experts(gm):
-        for node in gm.graph.nodes:
-            if node.op == "call_function" and "torch_moe" in str(node.target):
-                return len(node.args[3])  # w1_weight list length
-        return 0
-
-    assert _count_moe_experts(gm_full) == num_experts
-    assert _count_moe_experts(gm_reduced) == num_experts
+    assert _count_moe_experts_in_graph(gm_full) == num_experts
+    assert _count_moe_experts_in_graph(gm_reduced) == num_experts

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py`
around lines 331 - 335, Remove the duplicate nested function _count_moe_experts
inside the test (the local definition that iterates gm.graph.nodes and checks
for "torch_moe" nodes), and replace its usages with the module-level
_count_moe_experts_in_graph; i.e., delete the nested def _count_moe_experts(gm)
and update any calls in the test to call _count_moe_experts_in_graph(gm) instead
so the single shared implementation is used.

---

432-453: Consider adding numerical correctness checks to match the SimpleMoEForExport test pattern.

The test only validates graph structure and state-dict keys; it never calls the exported modules and asserts output equality. The corresponding test_moe_export_with_reduced_experts test (lines 337–349) asserts both all_close(y_full, y_reduced) and all_close(y_loaded, y_full). Without equivalent checks here, the GLM4 path could silently produce wrong outputs while still passing all structural assertions.

♻️ Proposed addition

     gm_reduced.load_state_dict(model.state_dict(), strict=False)
+
+    # Numerical: reduced export must produce the same output as full export
+    with torch.no_grad():
+        y_full = gm_full(**sample_kwargs)
+        y_reduced_after_load = gm_reduced(**sample_kwargs)
+    assert all_close(y_full.logits, y_reduced_after_load.logits), (
+        "Output after state-dict reload into reduced export differs from full export"
+    )

     # Source model must be fully restored

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py`
around lines 432 - 453, Add numerical output equality checks like in the
SimpleMoEForExport test: run a fixed input through gm_full and gm_reduced
(before and after calling gm_reduced.load_state_dict(model.state_dict(),
strict=False)), capture outputs (y_full, y_reduced, y_loaded) and assert they
match using the existing all_close or torch.allclose helper; keep the existing
structural/state-dict assertions but append asserts all_close(y_full, y_reduced)
and all_close(y_loaded, y_full) to ensure the GLM4 export path produces
numerically equivalent outputs. Use the existing symbols gm_full, gm_reduced,
model.state_dict(), and the helper all_close/torch.allclose to locate where to
add these checks.

tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py (4)

174-213: First two tests use >= 1 for fused op count; consider exact count for consistency.

test_fuse_add_cast_rms_norm and test_fuse_add_rms_norm_no_cast assert _count_fused_ops(gm_t) >= 1, while the later tests assert exact counts (== 1, == 2). Since each of these models has exactly one fusible pair, == 1 would be a tighter and more informative assertion.

Proposed fix

-    assert _count_fused_ops(gm_t) >= 1, "fused op not found in graph"
+    assert _count_fused_ops(gm_t) == 1, "Expected exactly 1 fused op in graph"

(Apply to both test_fuse_add_cast_rms_norm and test_fuse_add_rms_norm_no_cast.)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py`
around lines 174 - 213, Both tests use a loose assertion "_count_fused_ops(gm_t)
>= 1"; change this to an exact check "== 1" in test_fuse_add_cast_rms_norm and
test_fuse_add_rms_norm_no_cast to match later tests and reflect that each model
produces exactly one fused op; update the assertion strings if desired (e.g.,
keep "fused op not found in graph") so the checks in those functions
(_count_fused_ops and _count_rms_norm_ops called after _apply_transform on gm)
are consistent and stricter.

---

11-12: Import style: prefer importing the module, not the class directly.

Per coding guidelines, use from package.subpackage import module rather than from package.subpackage.module import ClassName. Consider:

from tensorrt_llm._torch.auto_deploy.transform import interface
from tensorrt_llm._torch.auto_deploy.transform.library import fused_add_rms_norm

Then reference interface.TransformConfig and fused_add_rms_norm.FuseAddRMSNorm. Note: the same pattern exists on unchanged lines (9–10) so this may warrant a broader cleanup. As per coding guidelines: "Python imports must use from package.subpackage import module style; never use from module import Class."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py`
around lines 11 - 12, The tests import classes directly; change to module-level
imports to follow project style: replace direct class imports of TransformConfig
and FuseAddRMSNorm with module imports (import the transform.interface module
and the transform.library.fused_add_rms_norm module) and update usages to
reference interface.TransformConfig and fused_add_rms_norm.FuseAddRMSNorm so all
references to TransformConfig and FuseAddRMSNorm use the module-qualified names.

---

138-140: _count_add_ops is defined but never used.

This helper is not called anywhere in the file. Consider removing it or adding a test assertion that uses it (e.g., verifying add ops are eliminated after fusion).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py`
around lines 138 - 140, The helper function _count_add_ops is defined but
unused; either remove the function or use it to assert fusion behavior—update
the test (e.g., in test_fused_add_rms_norm) to call _count_add_ops(gm) before
and after running the transformation pass and add an assertion that the number
of torch.ops.aten.add.Tensor nodes decreases (or becomes zero) after fusion, or
simply delete the _count_add_ops definition if no such assertion is desired;
reference the function name _count_add_ops and the graph module variable gm when
making the change.

---

6-8: Remove unused # noqa directive.

Ruff (RUF100) flags the blanket # noqa on line 6 as unused. The import is actively used (in _count_fused_ops), so no suppression is needed.

Proposed fix

-from tensorrt_llm._torch.auto_deploy.custom_ops.normalization.flashinfer_fused_add_rms_norm import (  # noqa
-    flashinfer_fused_add_rms_norm,
-)
+from tensorrt_llm._torch.auto_deploy.custom_ops.normalization.flashinfer_fused_add_rms_norm import (
+    flashinfer_fused_add_rms_norm,
+)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py`
around lines 6 - 8, Remove the unnecessary blanket "# noqa" on the import line
that brings in flashinfer_fused_add_rms_norm; the import is used (e.g., by
_count_fused_ops) so simply delete the "# noqa" directive from the import
statement importing flashinfer_fused_add_rms_norm to satisfy Ruff RUF100.

tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py (1)

14-15: Consider using built-in generic types instead of typing imports.

List, Tuple, and Optional from typing can be replaced with the built-in equivalents (list, tuple, X | None) since the codebase targets Python 3.10+.

♻️ Proposed refactor

-from typing import List, Optional, Tuple

-        matches: List[Tuple[Node, Optional[Node], Node]] = []
+        matches: list[tuple[Node, Node | None, Node]] = []

-            cast_node: Optional[Node] = None
+            cast_node: Node | None = None

And the return type annotation on _apply:

-    ) -> Tuple[GraphModule, TransformInfo]:
+    ) -> tuple[GraphModule, TransformInfo]:

Based on learnings: TensorRT-LLM requires Python ≥ 3.10, so PEP 585 built-in generics and X | None union syntax are available throughout the codebase.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py`
around lines 14 - 15, Replace typing generics with Python 3.10+ built-ins:
remove the imports of List, Optional, Tuple and use list, tuple, and X | None in
annotations instead. Update the module-level import statement (remove
List/Optional/Tuple) and change any function/type hints that use List[T],
Tuple[...], or Optional[X] to list[T], tuple[...], and X | None respectively;
also update the return type annotation on the method/function named _apply to
use the built-in generics/union syntax.

tensorrt_llm/_torch/auto_deploy/utils/_graph.py (1)

23-33: Module-level caches for derived op infrastructure look good.

The lazy-creation pattern with _derived_op_libs and _derived_op_registry keeps registrations persistent and avoids redundant work. One minor naming nit per coding guidelines: module-level mutable globals are conventionally G_-prefixed with upper snake_case (e.g., G_DERIVED_OP_LIBS), though the leading underscore already signals these are internal.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/auto_deploy/utils/_graph.py` around lines 23 - 33, Rename
the module-level mutable globals to follow the project's naming convention:
change _derived_op_libs to G_DERIVED_OP_LIBS and _derived_op_registry to
G_DERIVED_OP_REGISTRY, update all references/usages throughout the module
(including any initialization and lookups) to the new names, and preserve their
types and initial values (Dict[str, torch.library.Library] and Dict[str,
Callable]) so behavior is unchanged; ensure any import or external references in
this file are also updated to the new identifiers.

tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_create_derived_custom_op.py (1)

1-7: Comprehensive test suite for create_derived_custom_op.

Good coverage of correctness, registration visibility, caching semantics, fake implementation (default and custom), schema preservation with defaults, and OpOverload input handling.

Minor: the import on Line 6 imports create_derived_custom_op directly from _graph. Per coding guidelines, prefer from tensorrt_llm._torch.auto_deploy.utils import _graph and reference as _graph.create_derived_custom_op. This is a nit though.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_create_derived_custom_op.py`
around lines 1 - 7, The test imports create_derived_custom_op directly from the
_graph module; change the import to bring in the module instead and reference
the function via the module to follow the coding guideline: replace the direct
import with importing tensorrt_llm._torch.auto_deploy.utils import _graph and
update all usages of create_derived_custom_op in the test to call
_graph.create_derived_custom_op (ensure any other references in this test file
are updated accordingly).

tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.py (2)

27-30: Import style violates coding guidelines.

The guideline requires from package.subpackage import module rather than importing individual symbols from a module. This should import the module and access its members via the module namespace.

Proposed fix

-from tensorrt_llm._torch.auto_deploy.custom_ops.attention.trtllm_attention import (
-    _GlobalTrtllmPlanner,
-    prepare_trtllm_metadata_host,
-)
+from tensorrt_llm._torch.auto_deploy.custom_ops.attention import trtllm_attention

Then use trtllm_attention._GlobalTrtllmPlanner and trtllm_attention.prepare_trtllm_metadata_host throughout.

As per coding guidelines: "Python imports must use from package.subpackage import module style; never use from module import Class".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.py`
around lines 27 - 30, Replace the symbol-level import of _GlobalTrtllmPlanner
and prepare_trtllm_metadata_host with a module import so code follows the "from
package.subpackage import module" guideline: import the module
tensorrt_llm._torch.auto_deploy.custom_ops.attention.trtllm_attention as
trtllm_attention, then update all usages to call
trtllm_attention._GlobalTrtllmPlanner and
trtllm_attention.prepare_trtllm_metadata_host so references resolve via the
module namespace.

---

462-581: Unused loop variables and minor linting issues in the paged KV-cache test.

Lines 513 and 564: s is captured in for i, s in enumerate(seq_lengths) but never used in the loop body. Rename to _ to signal intent.

Proposed fix

-    for i, s in enumerate(seq_lengths):
+    for i, _ in enumerate(seq_lengths):

(Apply at both line 513 and line 564.)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.py`
around lines 462 - 581, In test_trtllm_attention_with_paged_kvcache, there are
two for-loops using unused loop variable s (for i, s in enumerate(seq_lengths))
around the prefill page assignment and the decode page-update sections; rename s
to _ in both loops so it signals an intentionally unused variable and fixes the
lint warnings (update the two occurrences of "for i, s in
enumerate(seq_lengths)" to "for i, _ in enumerate(seq_lengths)" while leaving
surrounding logic, page_assignments, pages_per_seq_list, and cache_locs handling
unchanged).

tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py (2)

14-31: Import style: several imports pull specific symbols from modules.

Per the coding guidelines, Python imports should use from package.subpackage import module style. Multiple imports here import specific classes/functions directly (e.g., from torch.fx import GraphModule, Node, from ...utils._graph import create_derived_custom_op). Consider importing the modules and using qualified access.

As per coding guidelines: "Python imports must use from package.subpackage import module style; never use from module import Class".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py`
around lines 14 - 31, Replace all direct symbol imports with module-level
imports and update usages to qualified names: import torch.fx as fx and replace
GraphModule/Node with fx.GraphModule and fx.Node; import ...utils._graph as
_graph and use _graph.create_derived_custom_op; import
...utils.multi_stream_utils as multi_stream_utils and use
multi_stream_utils._make_aux_stream_impl,
multi_stream_utils.cuda_stream_manager, and
multi_stream_utils.record_event_passthrough; import ...utils.node_utils as
node_utils and use node_utils.is_op; import ...shim.interface as shim_interface
and use shim_interface.CachedSequenceInterface; import ...models.factory as
factory and use factory.ModelFactory; import ..interface as interface and use
interface.BaseTransform, interface.SharedConfig, interface.TransformInfo,
interface.TransformRegistry. Update all references in this file accordingly so
no `from ... import Symbol` remains.

---

126-185: Graph mutation during iteration over pre-computed pairs — verify safety.

pairs is computed before the loop (line 144), and each iteration modifies the graph by inserting rec_node, replacing kv_linear with new_node, and erasing kv_linear. While the current logic appears safe (fork_point remains valid, Q-linears are correctly re-queried from live fork_point.users, and node_order is only used for original Q-linear nodes), this pattern is subtle.

One edge case: if two fork points share a Q-linear user (unlikely in MLA but possible in unusual graphs), the node_order dict built at line 149 could be stale for the second fork point's earliest_q lookup. Consider adding a brief comment about the safety invariant (e.g., "pairs are independent: no fork_point shares Q-linears with another").

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py`
around lines 126 - 185, The loop in _execute_kv_proj_in_aux_stream builds
node_order once from graph.nodes but then mutates the graph (inserting rec_node,
replacing/erasing kv_linear), which can make the precomputed node_order stale
when locating earliest_q for later fork_point iterations; fix by either
recomputing node_order (i.e., node_order = {n: i for i, n in
enumerate(graph.nodes)}) immediately before computing q_linears/earliest_q
inside the for fork_point, kv_linear in pairs loop, or explicitly document and
assert the safety invariant (e.g., add a clear comment and an assert in
_execute_kv_proj_in_aux_stream that pairs are independent: no two fork_point
entries share Q-linear users) so future maintainers know why the single
precomputed node_order is safe.

tensorrt_llm/_torch/auto_deploy/config/default.yaml (1)

142-143: Inconsistent enabled defaults between fuse_nvfp4_swiglu and fuse_swiglu.

fuse_nvfp4_swiglu (line 142) is implicitly enabled (no enabled: false), while fuse_swiglu (line 167) is explicitly disabled. Both depend on their respective pattern matchers (match_nvfp4_swiglu_pattern and match_swiglu_pattern) which are disabled by default, so this is safe in practice — but the inconsistency is worth noting. If both fusors are intended to be no-ops when their matchers haven't run, consider aligning their enabled defaults for clarity.

Also applies to: 167-169

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/auto_deploy/config/default.yaml` around lines 142 - 143,
The two fusor entries are inconsistent: fuse_nvfp4_swiglu lacks an explicit
enabled flag while fuse_swiglu is explicitly disabled; both depend on
match_nvfp4_swiglu_pattern and match_swiglu_pattern which are disabled by
default. Make the defaults explicit and consistent—e.g., add enabled: false to
fuse_nvfp4_swiglu (or alternatively set enabled: true on fuse_swiglu if you
intend both to be active) so fuse_nvfp4_swiglu and fuse_swiglu share the same
default behavior and avoid confusion.

tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py (3)

30-31: Use collections.deque for BFS queues instead of list.pop(0).

list.pop(0) is O(n) because it shifts every remaining element. In _find_merge_add (line 31) and _execute_shared_expert_in_aux_stream (line 114) the queues are used for BFS. Using collections.deque with popleft() gives O(1) dequeue.

Proposed fix for _find_merge_add

+from collections import deque
+
 def _find_merge_add(moe_node: Node) -> Optional[Node]:
     visited: Set[Node] = set()
-    queue = list(moe_node.users.keys())
+    queue = deque(moe_node.users.keys())
     while queue:
-        n = queue.pop(0)
+        n = queue.popleft()
         if n in visited:
             continue
         visited.add(n)
         if is_op(n, torch.ops.aten.add.Tensor):
             return n
         queue.extend(n.users.keys())
     return None

Apply the same change to the BFS in _execute_shared_expert_in_aux_stream (line 111–131).

Also applies to: 113-114

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py` around
lines 30 - 31, The BFS loops in _find_merge_add and
_execute_shared_expert_in_aux_stream use a Python list as "queue" and call
pop(0), which is O(n); change the queue to a collections.deque and replace
pop(0) with popleft() for O(1) dequeues. Import deque from collections (or use
collections.deque), initialize queue = deque(initial_items) where the list was
created, and update all usages in those functions (including any while queue: n
= queue.pop(0)) to n = queue.popleft(); ensure any code that extends the queue
uses append/extend or appendleft as appropriate. Make no other logic changes.

---

95-102: Potential misclassification when both add inputs descend from the MoE node.

If both arg0 and arg1 are in moe_node's ancestor set (e.g., they share an upstream path through the MoE node), the routed/shared classification would be wrong. The current check only tests arg0_ancestors; if moe_node is not in arg0_ancestors it defaults to assigning arg1 as routed, without verifying arg1 actually descends from moe_node. In practice the warning at line 133 would catch the failure later, but a symmetric check would be more robust.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py` around
lines 95 - 102, The current logic only checks _get_ancestors(arg0) to decide
which add input is routed vs shared, which can misclassify when both arg0 and
arg1 descend from moe_node; compute arg1_ancestors = _get_ancestors(arg1) and
then choose routed_output/shared_output with a symmetric decision: if moe_node
in arg0_ancestors or arg0 is moe_node set routed=arg0/shared=arg1; elif moe_node
in arg1_ancestors or arg1 is moe_node set routed=arg1/shared=arg0; otherwise
keep the existing fallback and warning behavior; update the block around
add_node, arg0, arg1, _get_ancestors, routed_output, and shared_output
accordingly.

---

153-168: Use replace_input_with() to safely handle both args and kwargs when rewiring fork_point.

The rewiring at lines 166–168 only substitutes fork_point in first_shared.args. Since line 146 checks all_input_nodes (which includes both args and kwargs), if fork_point appears in kwargs, the substitution will be missed, silently leaving a broken dependency. Use node.replace_input_with(fork_point, begin_aux_node) instead—it's available since PyTorch 1.11+ and already used throughout the codebase in similar transforms (e.g., sharding.py, gather_last_token_ids.py). Apply the same fix to lines 178–180 for add_node.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py` around
lines 153 - 168, The current rewiring replaces fork_point only inside
first_shared.args (and similarly for add_node.args), which misses occurrences in
kwargs; replace that manual tuple rebuild with the PyTorch helper
node.replace_input_with(fork_point, begin_aux_node) so both args and kwargs are
updated safely. Locate where you create begin_aux_node via
graph.call_function(begin_aux_stream_passthrough, args=(fork_point,)) and then
call first_shared.replace_input_with(fork_point, begin_aux_node); do the same
for add_node (replace its fork_point inputs with begin_aux_node) to ensure all
inputs/kwargs are patched.

tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_swiglu.py (1)

5-5: Avoid wildcard imports; import the module instead.

from ... import * violates the project's import convention and makes it unclear which names are being imported. Since this import is used for its side effect (registering custom ops), import the module itself:

Proposed fix

-from tensorrt_llm._torch.auto_deploy.custom_ops.linear.swiglu import *  # noqa
+from tensorrt_llm._torch.auto_deploy.custom_ops.linear import swiglu as _swiglu_ops  # noqa: F401 — side-effect import registers ops

As per coding guidelines: "Python imports must use from package.subpackage import module style; never use from module import Class."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_swiglu.py`
at line 5, Replace the wildcard import with a module import so names are not
leaked; change the line that currently does "from
tensorrt_llm._torch.auto_deploy.custom_ops.linear.swiglu import *" to "from
tensorrt_llm._torch.auto_deploy.custom_ops.linear import swiglu" (or import the
module under an alias) so the swiglu module is imported for its side effects
while avoiding a star import.

tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_moe.py (1)

96-102: Consider using torch.testing.assert_close for consistent assertion style.

_assert_numerical_correctness and _assert_cuda_graph_correctness use torch.allclose with a manual error message. torch.testing.assert_close provides richer diagnostics (element-wise diff summaries) and is used elsewhere in the test codebase (e.g., test_fuse_swiglu.py). A minor consistency nit.

Also applies to: 105-126

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_moe.py`
around lines 96 - 102, Replace the custom torch.allclose assertion in
_assert_numerical_correctness (and similarly in _assert_cuda_graph_correctness)
with torch.testing.assert_close to get richer, consistent diagnostics; call
torch.testing.assert_close(out, ref, atol=atol) (or include rtol if needed)
instead of the current assert torch.allclose(...) and remove the manual max-diff
formatting, keeping the same inputs gm, model, and test_x so the outputs being
compared remain identical.